"""
简单的SSH WebSocket Consumer - 参考Luna实现
"""
import json
import paramiko
import logging
import threading
import time
import select
import socket
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth.models import User
from .models import Host
from .jumpserver_audit import jumpserver_audit

logger = logging.getLogger(__name__)


class SimpleSSHConsumer(AsyncWebsocketConsumer):
    """简单的SSH WebSocket Consumer"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.host_id = None
        self.host = None
        self.user = None
        self.is_connected = False
        self.terminal_size = {'cols': 80, 'rows': 24}
        self.audit_session = None  # JumpServer审计会话
        self.command_buffer = ""   # 命令缓冲区

    async def connect(self):
        """WebSocket连接建立"""
        try:
            # 获取用户信息
            self.user = self.scope.get("user")
            # logger.info(f"WebSocket连接尝试，用户: {self.user}")  # 注释掉避免终端输出干扰

            if not self.user or not getattr(self.user, 'is_authenticated', False):
                logger.warning("WebSocket连接被拒绝：用户未认证")
                await self.close(code=4001)
                return

            # 获取主机ID
            self.host_id = self.scope['url_route']['kwargs']['host_id']
            # logger.info(f"WebSocket连接主机ID: {self.host_id}")  # 注释掉避免终端输出干扰

            # 接受WebSocket连接
            await self.accept()
            # logger.info("WebSocket连接已接受")  # 注释掉避免终端输出干扰

        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}")
            await self.close(code=4000)

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        await self.cleanup_ssh_connection()

    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')

            if message_type == 'auth':
                # 处理认证和连接
                ssh_port = data.get('ssh_port', 22)
                ssh_username = data.get('ssh_username')
                ssh_password = data.get('ssh_password')
                terminal_size = data.get('terminal_size', {'cols': 80, 'rows': 24})
                environment = data.get('environment', {})

                self.terminal_size = terminal_size
                self.environment = environment
                self.ssh_port = ssh_port
                self.ssh_username = ssh_username
                self.ssh_password = ssh_password

                await self.establish_ssh_connection()

            elif message_type == 'connect':
                # 建立SSH连接（兼容旧版本）
                await self.establish_ssh_connection()

            elif message_type == 'input' and self.ssh_channel:
                # 发送用户输入
                input_data = data.get('data', '')
                await self.send_input_to_ssh(input_data)

            elif message_type == 'resize':
                # 调整终端大小
                cols = data.get('cols', 80)
                rows = data.get('rows', 24)
                await self.resize_terminal(cols, rows)

            elif message_type == 'disconnect':
                # 断开SSH连接
                await self.cleanup_ssh_connection()

            elif message_type == 'get_pwd':
                # 获取当前工作目录
                await self.get_current_working_directory()

        except json.JSONDecodeError:
            logger.error("接收到无效的JSON数据")
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")

    async def establish_ssh_connection(self):
        """建立SSH连接"""
        try:
            # 获取主机信息
            self.host = await self.get_host(self.host_id)
            if not self.host:
                await self.send_error("主机不存在")
                return

            # 在线程中建立SSH连接
            import asyncio
            await asyncio.get_event_loop().run_in_executor(
                None, self._create_ssh_connection
            )

            if self.is_connected:
                # logger.info(f"🔍 SSH连接成功，开始集成JumpServer审计...")  # 注释掉避免终端输出干扰

                # 创建JumpServer审计会话
                await self.create_audit_session()

                # 启动输出读取线程
                self.read_thread = threading.Thread(target=self._read_ssh_output)
                self.read_thread.daemon = True
                self.read_thread.start()

                await self.send(text_data=json.dumps({
                    'type': 'connected',
                    'status': 'success',
                    'message': f'成功连接到 {self.host.hostname}'
                }))
                # logger.info(f"✅ SSH连接和JumpServer审计集成完成")  # 注释掉避免终端输出干扰
            else:
                logger.error(f"❌ SSH连接失败")
                await self.send_error("SSH连接失败")

        except Exception as e:
            logger.error(f"建立SSH连接失败: {e}")
            await self.send_error(f"连接失败: {str(e)}")

    def _create_ssh_connection(self):
        """在线程中创建SSH连接"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # 处理localhost特殊情况
            hostname = self.host.public_ip
            if hostname in ['localhost', '127.0.0.1', '::1']:
                hostname = '127.0.0.1'

            logger.info(f"尝试SSH连接: {hostname}:{self.host.port}")

            # 根据认证方式连接
            try:
                if hasattr(self, 'ssh_username') and hasattr(self, 'ssh_password'):
                    # 使用前端传来的认证信息
                    logger.info(f"使用前端认证信息连接: 用户={self.ssh_username}, 端口={self.ssh_port}")
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.ssh_port,
                        username=self.ssh_username,
                        password=self.ssh_password,
                        timeout=15,  # 增加超时时间
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30  # 增加banner超时时间
                    )
                elif self.host.auth_type == 'password':
                    logger.info(f"使用密码认证连接: 用户={self.host.username}, 端口={self.host.port}")
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.host.port,
                        username=self.host.username,
                        password=self.host.password,
                        timeout=15,
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30
                    )
                elif self.host.auth_type == 'key':
                    logger.info(f"使用密钥认证连接: 用户={self.host.username}, 端口={self.host.port}")
                    from io import StringIO
                    private_key = paramiko.RSAKey.from_private_key(
                        StringIO(self.host.private_key),
                        password=self.host.passphrase if self.host.passphrase else None
                    )
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.host.port,
                        username=self.host.username,
                        pkey=private_key,
                        timeout=15,
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30
                    )
                else:
                    raise Exception("不支持的认证方式")

            except paramiko.ssh_exception.SSHException as e:
                error_msg = str(e)
                if "Error reading SSH protocol banner" in error_msg:
                    raise Exception(f"SSH连接被拒绝 - 可能原因：1) SSH服务未运行 2) 端口{self.host.port}被防火墙阻止 3) 主机{hostname}不可达。详细错误: {error_msg}")
                elif "Connection reset by peer" in error_msg:
                    raise Exception(f"连接被目标主机重置 - 可能原因：1) SSH服务配置拒绝连接 2) 防火墙规则 3) 网络问题。请检查主机{hostname}:{self.host.port}的SSH服务状态")
                else:
                    raise Exception(f"SSH协议错误: {error_msg}")
            except socket.timeout:
                raise Exception(f"连接超时 - 主机{hostname}:{self.host.port}无响应，请检查网络连接和防火墙设置")
            except socket.error as e:
                if e.errno == 111:  # Connection refused
                    raise Exception(f"连接被拒绝 - 主机{hostname}端口{self.host.port}未开放或SSH服务未运行")
                elif e.errno == 113:  # No route to host
                    raise Exception(f"无法到达主机{hostname} - 请检查网络路由和主机状态")
                else:
                    raise Exception(f"网络连接错误: {e}")
            except Exception as e:
                if "Authentication failed" in str(e):
                    raise Exception(f"认证失败 - 请检查用户名和密码是否正确")
                else:
                    raise e

            # 创建交互式shell - 使用前端传来的环境变量
            environment = getattr(self, 'environment', {})
            if not environment:
                # 如果没有前端环境变量，使用默认值
                environment = {
                    'TERM': 'xterm-256color',
                    'LANG': 'en_US.UTF-8',
                    'LC_ALL': 'en_US.UTF-8',
                    'COLUMNS': str(self.terminal_size['cols']),
                    'LINES': str(self.terminal_size['rows']),
                    'LS_COLORS': 'rs=0:di=01;34:ln=01;36:mh=00:pi=40;33:so=01;35:do=01;35:bd=40;33;01:cd=40;33;01:or=40;31;01:mi=00:su=37;41:sg=30;43:ca=30;41:tw=30;42:ow=34;42:st=37;44:ex=01;32',
                    'COLORTERM': 'truecolor',
                    'CLICOLOR': '1',
                    'CLICOLOR_FORCE': '1',
                    'LSCOLORS': 'ExFxCxDxBxegedabagacad'
                }

            self.ssh_channel = self.ssh_client.invoke_shell(
                term='xterm-256color',
                width=self.terminal_size['cols'],
                height=self.terminal_size['rows'],
                environment=environment
            )

            # 先通过SFTP预创建vim配置文件，避免在终端显示命令
            try:
                sftp = self.ssh_client.open_sftp()
                vim_config = """syntax on
set t_Co=256
colorscheme default
set background=dark
filetype plugin indent on
set nocompatible
"""
                with sftp.file('.vimrc', 'w') as f:
                    f.write(vim_config)
                sftp.close()
                logger.info("vim配置文件创建成功")
            except Exception as e:
                logger.warning(f"创建vim配置文件失败: {e}")

            # 等待shell初始化
            time.sleep(0.3)

            # 不发送欢迎信息，让用户看到原始的shell提示符
            # 这样避免了欢迎信息被当作命令执行的问题
            # 只发送一个回车来确保显示提示符
            self.ssh_channel.send(b'\r')
            time.sleep(0.2)

            self.is_connected = True
            logger.info(f"SSH连接成功: {self.host.hostname}")

        except Exception as e:
            logger.error(f"SSH连接失败: {e}")
            self.is_connected = False
            raise e



    def _read_ssh_output(self):
        """读取SSH输出的线程函数"""
        while self.is_connected and self.ssh_channel:
            try:
                # 使用select检查是否有数据可读
                ready, _, _ = select.select([self.ssh_channel], [], [], 0.1)

                if ready:
                    try:
                        data = self.ssh_channel.recv(65535)  # 增大缓冲区
                        if data:
                            # 普通文本数据处理
                            output = data.decode('utf-8', errors='replace')
                            self._send_output_to_websocket(output)
                        else:
                            # 连接关闭
                            break
                    except socket.timeout:
                        continue
                    except Exception as e:
                        logger.error(f"接收SSH数据失败: {e}")
                        break

            except Exception as e:
                logger.error(f"SSH输出读取失败: {e}")
                break

        logger.info("SSH输出读取线程结束")

    def _send_output_to_websocket(self, output):
        """发送输出到WebSocket"""
        try:
            import asyncio

            # 记录输出到审计系统
            self._log_output_data(output)

            message = json.dumps({
                'type': 'output',
                'data': output
            })

            # 使用线程安全的方式发送消息
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                loop.run_until_complete(self.send(text_data=message))
            finally:
                loop.close()

        except Exception as e:
            logger.error(f"发送输出到WebSocket失败: {e}")

    def _log_output_data(self, output):
        """记录输出数据到审计系统"""
        try:
            if self.audit_session:
                jumpserver_audit.log_replay_data(
                    self.audit_session.session_id, 'output', output,
                    None, self.terminal_size['cols'], self.terminal_size['rows']
                )
        except Exception as e:
            logger.error(f"记录输出数据失败: {e}")

    async def send_input_to_ssh(self, input_data):
        """发送输入到SSH"""
        try:
            if self.ssh_channel and self.is_connected:
                # 记录输入到审计系统
                await self.log_input_data(input_data)

                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.send, input_data.encode('utf-8')
                )
        except Exception as e:
            logger.error(f"发送输入到SSH失败: {e}")

    async def log_input_data(self, input_data):
        """记录输入数据到审计系统"""
        try:
            if self.audit_session:
                logger.debug(f"🔍 记录输入数据到JumpServer审计: {repr(input_data)}")

                # 记录回放数据
                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, jumpserver_audit.log_replay_data,
                    self.audit_session.session_id, 'input', input_data,
                    None, self.terminal_size['cols'], self.terminal_size['rows']
                )

                # 处理命令识别
                await self.process_command_input(input_data)
            else:
                logger.warning(f"⚠️ 审计会话不存在，无法记录输入数据")

        except Exception as e:
            logger.error(f"❌ 记录输入数据失败: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")

    async def process_command_input(self, input_data):
        """处理命令输入，识别完整命令 - 改进版本支持Tab补全"""
        try:
            # 处理特殊字符
            if input_data == '\t':  # Tab键补全
                logger.debug(f"🔍 检测到Tab键补全，当前缓冲区: {repr(self.command_buffer)}")
                # Tab键补全时，不清空缓冲区，等待用户确认或继续输入
                return
            elif input_data == '\x7f' or input_data == '\b':  # 退格键
                if self.command_buffer:
                    self.command_buffer = self.command_buffer[:-1]
                    logger.debug(f"🔍 退格键，当前缓冲区: {repr(self.command_buffer)}")
                return
            elif input_data.startswith('\x1b'):  # 转义序列（方向键等）
                # 处理方向键历史命令
                if input_data == '\x1b[A':  # 上箭头
                    logger.debug(f"🔍 检测到上箭头键（历史命令）")
                    # 清空当前缓冲区，因为用户可能选择了历史命令
                    self.command_buffer = ""
                elif input_data == '\x1b[B':  # 下箭头
                    logger.debug(f"🔍 检测到下箭头键（历史命令）")
                    self.command_buffer = ""
                else:
                    logger.debug(f"🔍 检测到其他转义序列: {repr(input_data)}")
                return

            # 检查是否是完整命令（以回车结尾）
            if '\r' in input_data or '\n' in input_data:
                # 获取当前行的完整内容（包括Tab补全后的内容）
                # 这里我们需要从SSH输出中提取实际执行的命令
                self._process_completed_command()

                # 清空命令缓冲区
                self.command_buffer = ""
                return

            # 添加普通字符到命令缓冲区
            # 过滤掉一些特殊字符，但保留可见字符
            if input_data.isprintable() or input_data in [' ']:
                self.command_buffer += input_data
                logger.debug(f"🔍 添加字符到缓冲区: {repr(input_data)}, 当前缓冲区: {repr(self.command_buffer)}")

        except Exception as e:
            logger.error(f"❌ 处理命令输入失败: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")

    def _process_completed_command(self):
        """处理完成的命令 - 从缓冲区提取并清理"""
        try:
            if not self.command_buffer:
                return

            # 清理命令
            command = self.command_buffer.strip()

            # 清理ANSI转义序列和控制字符
            import re
            command = re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', command)  # ANSI转义序列
            command = re.sub(r'[\x00-\x1f\x7f]', '', command)  # 控制字符
            command = re.sub(r'\s+', ' ', command)  # 多个空格合并为一个
            command = command.strip()

            if command and len(command) > 0:
                logger.info(f"🔍 检测到完整命令: {command}")

                # 异步记录命令到审计系统
                import asyncio
                import threading

                def log_command_sync():
                    try:
                        jumpserver_audit.log_command(
                            self.audit_session.session_id, command, 'input'
                        )
                        logger.info(f"✅ 命令已记录到JumpServer审计: {command}")
                    except Exception as e:
                        logger.error(f"❌ 记录命令失败: {e}")

                # 在线程中执行，避免阻塞
                threading.Thread(target=log_command_sync, daemon=True).start()

        except Exception as e:
            logger.error(f"❌ 处理完成命令失败: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")

    async def get_current_working_directory(self):
        """获取当前工作目录 - 直接使用PWD环境变量"""
        try:
            if self.ssh_channel and self.is_connected:
                # 直接发送PWD环境变量给前端，不执行命令
                # 在shell初始化时PWD就已经设置好了
                default_pwd = f'/home/{self.host.username}' if self.host.username != 'root' else '/root'

                # 发送当前工作目录给前端
                await self.send(text_data=json.dumps({
                    'type': 'pwd_response',
                    'pwd': 'CURRENT_DIR',  # 让后端API使用当前SSH会话的工作目录
                    'use_current_dir': True
                }))
                logger.info(f"使用当前SSH会话工作目录")

        except Exception as e:
            logger.error(f"获取当前工作目录失败: {e}")
            # 发生错误时使用默认家目录
            default_pwd = f'/home/{self.host.username}' if hasattr(self, 'host') and self.host.username != 'root' else '/root'
            await self.send(text_data=json.dumps({
                'type': 'pwd_response',
                'pwd': default_pwd,
                'error': str(e)
            }))

    async def handle_sz_command(self, command):
        """处理sz命令，提供文件下载"""
        try:
            # 解析sz命令，提取文件路径
            parts = command.split(' ', 1)
            if len(parts) < 2:
                await self.send(text_data=json.dumps({
                    'type': 'output',
                    'data': '\r\n错误: sz命令需要指定文件路径\r\n'
                }))
                return

            file_path = parts[1].strip()

            # 发送sz命令到SSH
            await self.send_input_to_ssh(command + '\r')

            # 发送文件下载提示给前端
            await self.send(text_data=json.dumps({
                'type': 'sz_download',
                'file_path': file_path,
                'message': f'检测到sz命令，准备下载文件: {file_path}'
            }))

        except Exception as e:
            logger.error(f"处理sz命令失败: {e}")
            await self.send(text_data=json.dumps({
                'type': 'output',
                'data': f'\r\n错误: sz命令处理失败: {str(e)}\r\n'
            }))

    async def resize_terminal(self, cols, rows):
        """调整终端大小"""
        try:
            self.terminal_size = {'cols': cols, 'rows': rows}
            if self.ssh_channel and self.is_connected:
                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.resize_pty, cols, rows
                )

                # 记录resize事件到审计系统
                if self.audit_session:
                    resize_data = json.dumps({'cols': cols, 'rows': rows})
                    await asyncio.get_event_loop().run_in_executor(
                        None, jumpserver_audit.log_replay_data,
                        self.audit_session.session_id, 'resize', resize_data,
                        None, cols, rows
                    )

        except Exception as e:
            logger.error(f"调整终端大小失败: {e}")

    async def cleanup_ssh_connection(self):
        """清理SSH连接"""
        try:
            self.is_connected = False

            # 关闭审计会话
            if self.audit_session:
                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, jumpserver_audit.close_session,
                    self.audit_session.session_id, 0
                )
                self.audit_session = None

            if self.ssh_channel:
                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.close
                )
                self.ssh_channel = None

            if self.ssh_client:
                import asyncio
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_client.close
                )
                self.ssh_client = None

        except Exception as e:
            logger.error(f"清理SSH连接失败: {e}")

    async def create_audit_session(self):
        """创建JumpServer审计会话"""
        try:
            logger.info("🔍 开始创建JumpServer审计会话...")
            client_ip = self.get_client_ip()
            user_agent = self.get_user_agent()

            logger.info(f"🔍 审计会话参数 - 用户: {self.user.username}, 主机: {self.host.hostname}, IP: {client_ip}")

            # 在线程中创建审计会话
            import asyncio
            self.audit_session = await asyncio.get_event_loop().run_in_executor(
                None, jumpserver_audit.create_session,
                self.user, self.host, client_ip, user_agent
            )

            logger.info(f"✅ JumpServer审计会话创建成功: {self.audit_session.session_id}")

        except Exception as e:
            logger.error(f"❌ 创建JumpServer审计会话失败: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")

    def get_client_ip(self):
        """获取客户端IP"""
        headers = dict(self.scope.get('headers', []))

        # 尝试从X-Forwarded-For获取真实IP
        forwarded_for = headers.get(b'x-forwarded-for')
        if forwarded_for:
            return forwarded_for.decode().split(',')[0].strip()

        # 尝试从X-Real-IP获取
        real_ip = headers.get(b'x-real-ip')
        if real_ip:
            return real_ip.decode()

        # 使用客户端地址
        client = self.scope.get('client', ['unknown', 0])
        return client[0] if client else 'unknown'

    def get_user_agent(self):
        """获取用户代理"""
        headers = dict(self.scope.get('headers', []))
        user_agent = headers.get(b'user-agent')
        return user_agent.decode() if user_agent else 'Unknown'

    async def send_error(self, message):
        """发送错误消息"""
        await self.send(text_data=json.dumps({
            'type': 'error',
            'status': 'error',
            'message': message
        }))

    @database_sync_to_async
    def get_host(self, host_id):
        try:
            return Host.objects.get(id=host_id)
        except Host.DoesNotExist:
            return None
